/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2020-2023. All rights reserved.
 * Description: hccl实现基类文件
 */

#include <atomic>
#include <chrono>
#include <thread>
#include <algorithm>
#include <numeric>
#include <unordered_set>
#include <sys/time.h>
#include "slog.h"
#include "slog_api.h"
#include "dlprof_function.h"
#include "dlrt_function.h"
#include "externalinput_pub.h"
#include "p2p_mgmt_pub.h"
#include "../command_handle.h"
#include "config.h"
#include "transport_heterog_event_roce_pub.h"
#include "transport_heterog_event_tcp_pub.h"
#include "transport_heterog_roce_pub.h"
#include "transport_roce.h"
#include "alg_remote_lookup.h"
#include "alg_coll_remote_update.h"
#include "alg_coll_remote_lookup.h"
#include "callback_thread_manager.h"
#include "transport_shm_event_pub.h"
#include "stream_active_manager.h"
#include "device_capacity.h"
#include "task_exception_handler_pub.h"
#include "adapter_rts.h"
#include "adapter_prof.h"
#include "hccl_impl_rpc.h"
#include "spin_mutex.h"

using namespace std;

namespace hccl {

enum TransferMemInfoIdx {
    TRANSFER_MEM_INFO_KEY_IDX = 0,
    TRANSFER_MEM_INFO_VALUE_IDX = 1,
    TRANSFER_MEM_INFO_RDMA_ENVELOPE_IDX = 2,
    TRANSFER_MEM_INFO_IDX_NUM = 3
};

HcclImplRpc::HcclImplRpc():HcclImplBase()
{};

HcclImplRpc::~HcclImplRpc()
{};


HcclResult HcclImplRpc::RegRemoteLookUpExecutor(s32 tag, bool isWorker, bool isKernelLaunch, HcclOpPtr opPtr)
{
    HCCL_INFO("transportHeterogMap_ size[%u]", transportHeterogMap_.Size());

    if (opPtr == nullptr) {
        auto it = tagOpExecutorMap_.Find(tag);
        if (it.second) {
            return HCCL_SUCCESS;
        }
    }

    map<u32, u32> psRankIdMap;
    if (isWorker) {
        CHK_RET(BuildEsWorkerTransportLinks(tag));
    }
    CHK_RET(GetPsMap(psRankIdMap));

    if (opPtr != nullptr) {
        HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
        ptr->param_.opExecutorPtr = new (nothrow) OpRemoteLookUpExecutor(mrManager_, dispatcher_, realUserRank_);
        CHK_PTR_NULL(ptr->param_.opExecutorPtr);
        CHK_RET(static_cast<OpRemoteLookUpExecutor *>(ptr->param_.opExecutorPtr)->Init(transportHeterogMap_[tag],
            psRankIdMap));
    } else {
        // 同tag当前无并发，故可先占位再申请，减少性能影响
        auto it = tagOpExecutorMap_.Emplace(tag, nullptr);

        unique_ptr<OpExecutor> &executor = it.first->second;
        EXECEPTION_CATCH((executor = make_unique<OpRemoteLookUpExecutor>(mrManager_, dispatcher_, realUserRank_)),
            return HCCL_E_PTR);

        OpRemoteLookUpExecutor *lookupExe = dynamic_cast<OpRemoteLookUpExecutor *>(executor.get());
        CHK_PTR_NULL(lookupExe);
        CHK_RET(lookupExe->Init(transportHeterogMap_[tag], psRankIdMap));
    }

    if (!isKernelLaunch) {
        return HCCL_SUCCESS;
    }

    CHK_RET(this->RegTransportsKernelLaunch(tag, psRankIdMap));

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::RegRemoteUpdateExecutor(s32 tag, bool isWorker, bool isKernelLaunch)
{
    HCCL_INFO("transportHeterogMap_ size[%u]", transportHeterogMap_.Size());

    auto it = tagOpExecutorMap_.Find(tag);
    if (it.second) {
        return HCCL_SUCCESS;
    }

    // key:psId value:rankId
    map<u32, u32> psRankIdMap;
    vector<u32> workerList;
    if (isWorker) {
        CHK_RET(BuildEsWorkerTransportLinks(tag));
    } else {
        CHK_RET(BuildEsPsTransportLinks(tag, false));
    }
    CHK_RET(GetPsMap(psRankIdMap));
    CHK_RET(GetWorkerList(workerList));

    // 同tag当前无并发，故可先占位再申请，减少性能影响
    auto exeIt = tagOpExecutorMap_.Emplace(tag, nullptr);

    unique_ptr<OpExecutor> &executor = exeIt.first->second;
    EXECEPTION_CATCH((executor = make_unique<OpCollRemoteUpdateExecutor>()), return HCCL_E_PTR);

    OpCollRemoteUpdateExecutor *updateExe = dynamic_cast<OpCollRemoteUpdateExecutor *>(executor.get());
    CHK_PTR_NULL(updateExe);

    bool needThread = !isWorker;
    CHK_RET(updateExe->Init(transportHeterogMap_[tag], psRankIdMap, workerList, embeddingReduceThreadNum_, needThread));

    if (!isKernelLaunch) {
        return HCCL_SUCCESS;
    }

    CHK_RET(this->RegTransportsKernelLaunch(tag, psRankIdMap));

    return HCCL_SUCCESS;
}

// RecoverValue下发到Aicore的unsorted_segment_sum算子进行执行
HcclResult HcclImplRpc::ValueReduceSumAicoreKernelLaunch(const HdcsEmbeddingServiceParam &para, void *stream)
{
// 当前在HCCD场景下不能使用
#ifndef HCCD
    // 可重复调用，内部会避免重复初始化
    CHK_RET(aiCoreUnsortedSegmentSum_.Init());

    void *mappingMatrix = para.pairedMode ? para.indices :
        (static_cast<s64 *>(para.keyTransferMem) + para.keyMaxNum);
    u32 valueDim = static_cast<u32>(para.valueItemSize) / SIZE_TABLE[para.valuesType];
    HcclDataType indicesType = HCCL_DATA_TYPE_INT32;

    CHK_RET(aiCoreUnsortedSegmentSum_.Run(para.values, mappingMatrix, para.keyMaxNum, para.keyMaxNum,
        valueDim, indicesType, stream, para.valueTransferMem));

    return HCCL_SUCCESS;
#else
    HCCL_ERROR("AiCore UnsortedSegmentSum not supported!");

    return HCCL_E_NOT_SUPPORT;
#endif
}

// ClearValue下发到Aicore的算子进行执行
HcclResult HcclImplRpc::ValueClearAicoreKernelLaunch(const HdcsEmbeddingServiceParam &para, void *stream)
{
// 当前在HCCD场景下不能使用
#ifndef HCCD
    // 可重复调用，内部会避免重复初始化
    CHK_RET(aiCoreCrackCleard_.Init());

    vector<int64_t> crackAddr;
        // 字节数量
    vector<int64_t> crackSize;
    crackAddr.push_back(reinterpret_cast<int64_t>(para.valueTransferMem));

    u64 size = para.keyMaxNum * static_cast<u32>(para.valueItemSize);
    CHK_PRT_RET(UNLIKELY(size > INT64_MAX),
        HCCL_ERROR("value size[%llu] > INT64_MAX, Aicore Cleard does not support", size), HCCL_E_PARA);

    crackSize.push_back(static_cast<int64_t>(size));

    CHK_RET(aiCoreCrackCleard_.Run(crackAddr, crackSize, stream));

    return HCCL_SUCCESS;
#else
    HCCL_ERROR("AiCore CrackCleared not supported!");

    return HCCL_E_NOT_SUPPORT;
#endif
}

HcclResult HcclImplRpc::CollRecvLookupRequest(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);

    std::unique_lock<std::mutex> lock(tagServiceHandleMapMutex_);
    tagServiceHandleMap_[ptr->param_.tag] = opPtr;
    lock.unlock();

    ptr->param_.getRequestParam.rankId = userRank_;
    ptr->param_.algPtr = new (std::nothrow) AlgCollRemoteLookUp();
    AlgCollRemoteLookUp *alg = static_cast<AlgCollRemoteLookUp *>(ptr->param_.algPtr);
    CHK_PTR_NULL(alg);
    CHK_RET(alg->Process(opPtr));

    OpRemoteLookUpExecutor *opExecutorPtr = static_cast<OpRemoteLookUpExecutor *>(ptr->param_.opExecutorPtr);
    CHK_PTR_NULL(opExecutorPtr);
    CHK_RET(static_cast<OpRemoteLookUpExecutor *>(ptr->param_.opExecutorPtr)->ExecuteOp(opPtr));

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::UpdateTaskSendByRdma(const HdcsEmbeddingServiceParam &para, void *stream)
{
    HcclEsRdmaInfoForUpdate *devRdmaEnveInfos = static_cast<HcclEsRdmaInfoForUpdate *>(para.rdmaEnveInfosTransferMem);
    unique_ptr<HcclEsRdmaInfoForUpdate[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMems_[para.tag];

    for (u32 i = 0; i < serversInfo_.size(); ++i) {
        u32 rankId = serversInfo_[i].rankId;
        u32 psId = i;

        if (rdmaEnveInfos[psId].commType != HcclHeterogCommType::RDMA) {
            HCCL_DEBUG("rankId[%u], psId[%u] send transport type is not rdma, commType[%d]",
                rankId, psId, rdmaEnveInfos[psId].commType);
            continue;
        }

        HCCL_DEBUG("rankId[%u], psId[%u] send transport type is rdma", rankId, psId);

        TransportHeterogRoce *transportRoce =
            dynamic_cast<TransportHeterogRoce *>(transportHeterogMap_[para.tag][rankId].get());
        CHK_PTR_NULL(transportRoce);

        HcclEnvelope &devEnvelope = devRdmaEnveInfos[psId].envelope;

        CHK_RET(transportRoce->SendEnvelope(devEnvelope, stream));

        HcclEnvelope &devEnvelopeValue = devRdmaEnveInfos[psId].envelopeValue;

        HCCL_DEBUG("Begin to SendEnvelopeValue");
        CHK_RET(transportRoce->SendEnvelope(devEnvelopeValue, stream));
    }

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::NotifyAndResetKernelLaunch(const HdcsEmbeddingServiceParam &para, void *stream)
{
    auto &profilingManager = hccl::ProfilingManager::Instance();

    unique_ptr<HcclEsRdmaInfoForUpdate[]> &rdmaEnveInfos = rdmaEnvelopeInfoHostMems_[para.tag];

    shared_ptr<RemoteNotify> isendDoneNotify = nullptr;
    // worker通知ps不使用notify，走rdma时只需要在worker侧等notify
    for (u32 i = 0; i < serversInfo_.size(); ++i) {
        u32 rankId = serversInfo_[i].rankId;

        if (rdmaEnveInfos[i].commType != HcclHeterogCommType::PCIE) {
            HCCL_DEBUG("rankId[%u], psId[%u] transport type[%d] is not pcie, no need to sendkey record",
                rankId, i, rdmaEnveInfos[i].commType);
            continue;
        }

        TransportShmEvent *transport =
            dynamic_cast<TransportShmEvent *>(transportHeterogMap_[para.tag][rankId].get());
        CHK_PTR_NULL(transport);
        CHK_RET(transport->GetIsendDoneSignal(isendDoneNotify));
        // isendDoneNotify key，通知对端用
        Stream tmpStream(stream);
        isendDoneNotify->Post(tmpStream, dispatcher_);
        // isendDoneNotify value
        isendDoneNotify->Post(tmpStream, dispatcher_);
    }

    // 当前无流水，且ResetUniqueHandle处理不影响PS的数据接收，故可以不必用remoteImrecvDoneSignal进行notify等待

    rtKernelLaunchNames_t launchNames;
    launchNames.soName = "libccl_kernel.so";
    launchNames.kernelName = "HdcsCollRemoteUpdateResetUniqueHandle";
    launchNames.opName = "CollRemoteUpdateResetUniqueHandle";
    // 清除去重信息，以及reduceSum的第二次value清零处理

    uint64_t beginTime = hrtMsprofSysCycleTime();
    EsUpdatePara esUpdatePara;
    CHK_RET(SetupEsUpdatePara(esUpdatePara, para));
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
        ProfTaskType::TASK_UPDATE_RESET_UNIQUE_HANDLE));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsEmbeddingServiceParam), nullptr, stream));
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime,
        ProfTaskType::TASK_UPDATE_RESET_UNIQUE_HANDLE));

    std::shared_ptr<LocalIpcNotify> remoteIsendDoneNotify = nullptr;
    for (u32 i = 0; i < serversInfo_.size(); ++i) {
        u32 rankId = serversInfo_[i].rankId;

        TransportHeterog *transport = transportHeterogMap_[para.tag][rankId].get();
        CHK_PTR_NULL(transport);

        // remoteIsendDoneNotify
        CHK_RET(transport->GetRemoteIsendDoneSignal(remoteIsendDoneNotify));

        Stream tmpStream(stream);
        uint64_t beginTime = hrtMsprofSysCycleTime();
        CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
            ProfTaskType::TASK_NOTIFY_REMOTE_IMRECV_DONE_SIGNAL_VALUE));

        HCCL_DEBUG("wait ps send complete");
        // worker侧等ps侧回复responce
        TransportHeterogRoce *tmpTransport = dynamic_cast<TransportHeterogRoce*>(transport);
        if (tmpTransport != nullptr && tmpTransport->tagQpInfo_.qpMode == OPBASE_QP_MODE_EXT) {
            auto meta = HcclOpMetaInfo::GetOneForRecieve();
            CHK_RET(InitTask(dispatcher_, tmpStream, meta.isEnableCache, meta.GetCacheKey()));
            remoteIsendDoneNotify->Wait(tmpStream, dispatcher_);
            CHK_RET(LaunchTask(dispatcher_, tmpStream));
        } else {
            remoteIsendDoneNotify->Wait(tmpStream, dispatcher_);
        }
        CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime,
            ProfTaskType::TASK_NOTIFY_REMOTE_IMRECV_DONE_SIGNAL_VALUE));
    }

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::CollIsendLookupResponse(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);

    auto it = tagRankTransMemsMap_.Find(ptr->param_.tag);
    if (it.second) {
        ptr->param_.getRequestParam.transferMems = it.first->second;
    }

    AlgCollRemoteLookUp *alg = static_cast<AlgCollRemoteLookUp *>(ptr->param_.algPtr);
    CHK_PTR_NULL(alg);
    CHK_RET(alg->Process(opPtr));
    OpRemoteLookUpExecutor *opExecutorPtr = static_cast<OpRemoteLookUpExecutor *>(ptr->param_.opExecutorPtr);
    CHK_PTR_NULL(opExecutorPtr);
    CHK_RET(opExecutorPtr->ExecuteOp(opPtr));
    delete alg;
    ptr->param_.algPtr = nullptr;
    delete opExecutorPtr;
    ptr->param_.opExecutorPtr = nullptr;

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::CollRecvUpdateRequest(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);

    EXECEPTION_CATCH({ auto &alg = GenerateEmbAlg<AlgCollRemoteUpdate>(ptr->param_.tag); \
        CHK_SMART_PTR_NULL(alg); \
        ptr->param_.getRequestParam.rankId = userRank_; \
        CHK_RET(alg->Process(opPtr)); \
        }, return HCCL_E_PTR);

    u32 workerSize = ptr->param_.getRequestParam.workerList.size();
    // 流水的循环是否结束判断逻辑
    while (std::count(ptr->param_.getRequestParam.rankTransKeysCompFlag.begin(),
        ptr->param_.getRequestParam.rankTransKeysCompFlag.end(), 0) != workerSize) {
        HCCL_INFO("getRequestParam ExecuteOp start tag[%d]", ptr->param_.tag);
        CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    }

    dynamic_cast<OpCollRemoteUpdateExecutor *>(tagOpExecutorMap_[ptr->param_.tag].get())->ClearAllTask(ptr);
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::CollRemoteLookUp(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    auto &embeddingParam = ptr->param_.embeddingParam;

    u32 keyBytesNum = embeddingParam.enableKeyCounter ? ES_KEY_AND_COUNTER_MEM_BYTES_SIZE :
        SIZE_TABLE[ptr->param_.keyType];
    unique_lock<std::mutex> lockTransferMemsMap(transferMemsMapMutex_);
    const vector<DeviceMem> &memInfos = transferMemsMap_[ptr->param_.tag];
    lockTransferMemsMap.unlock();

    if (memInfos.size() != TRANSFER_MEM_INFO_IDX_NUM ||
        ptr->param_.keyMaxNum * keyBytesNum > memInfos[TRANSFER_MEM_INFO_KEY_IDX].size() ||
        ptr->param_.keyMaxNum * embeddingParam.valueItemSize >
        memInfos[TRANSFER_MEM_INFO_VALUE_IDX].size()) {
        HCCL_ERROR("[Remote][LookUp]transferMems invalid or too big");
        return HCCL_E_PARA;
    }

    embeddingParam.keyTransferMem = memInfos[TRANSFER_MEM_INFO_KEY_IDX].ptr();
    embeddingParam.keyTransferMemSize = memInfos[TRANSFER_MEM_INFO_KEY_IDX].size();
    embeddingParam.valueTransferMem = memInfos[TRANSFER_MEM_INFO_VALUE_IDX].ptr();
    embeddingParam.valueTransferMemSize = memInfos[TRANSFER_MEM_INFO_VALUE_IDX].size();

    embeddingParam.rdmaEnveInfosTransferMem = memInfos[TRANSFER_MEM_INFO_RDMA_ENVELOPE_IDX].ptr();
    embeddingParam.rdmaEnveInfosTransferMemSize = memInfos[TRANSFER_MEM_INFO_RDMA_ENVELOPE_IDX].size();

    embeddingParam.workerspaceMemSize = embeddingParam.keyTransferMemSize + embeddingParam.valueTransferMemSize +
        embeddingParam.rdmaEnveInfosTransferMemSize;

    HCCL_INFO("[Remote][LookUp]keyTransferMemSize[%llu] valueTransferMemSize[%llu] rdmaEnveInfosTransferMemSize[%llu]"
        " workerspaceMemSize[%llu] keyMaxNum[%llu] valueItemSize[%d]", embeddingParam.keyTransferMemSize,
        embeddingParam.valueTransferMemSize, embeddingParam.rdmaEnveInfosTransferMemSize,
        embeddingParam.workerspaceMemSize, ptr->param_.keyMaxNum, embeddingParam.valueItemSize);

    EXECEPTION_CATCH({ auto &alg = GenerateEmbAlg<AlgRemoteLookUp>(ptr->param_.tag); \
        CHK_SMART_PTR_NULL(alg); \
        CHK_RET(alg->Process(opPtr)); \
        }, return HCCL_E_PTR);

    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    CHK_RET(notifyPool_->UnregisterOp(ptr->param_.sTag));

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::CollRemoteLookUpPaired(HcclOpPtr opPtr)
{
    return CollRemoteLookUp(opPtr);
}

HcclResult HcclImplRpc::CollRemoteUpdate(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    auto &embeddingParam = ptr->param_.embeddingParam;

    EXECEPTION_CATCH({ auto &alg = GenerateEmbAlg<AlgCollRemoteUpdate>(ptr->param_.tag); \
        CHK_SMART_PTR_NULL(alg); \
        CHK_RET(alg->Process(opPtr)); \
        }, return HCCL_E_PTR);

    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));

    HdcsEmbeddingServiceParam para{};
    unique_lock<std::mutex> lockTransferMemsMap(transferMemsMapMutex_);
    const std::vector<DeviceMem> &memInfos = transferMemsMap_[ptr->param_.tag];
    lockTransferMemsMap.unlock();
    if (memInfos.size() != TRANSFER_MEM_INFO_IDX_NUM ||
        ptr->param_.keyMaxNum * embeddingParam.valueItemSize >
        memInfos[TRANSFER_MEM_INFO_VALUE_IDX].size()) {
        HCCL_ERROR("[Coll][RemoteUpdate]transferMems invalid or too big");
        return HCCL_E_PARA;
    }

    para.tag = ptr->param_.tag;
    para.keys = ptr->param_.keys;
    para.keyMaxNum = ptr->param_.keyMaxNum;
    para.keyType = HCCL_DATA_TYPE_INT64;
    para.values = ptr->param_.values;
    para.valueCount = ptr->param_.valueCount;
    para.valuesType = HCCL_DATA_TYPE_FP32;
    para.insertFlag = embeddingParam.insertFlag;
    para.valueItemSize = embeddingParam.valueItemSize;
    para.keyTransferMem = memInfos[TRANSFER_MEM_INFO_KEY_IDX].ptr();
    para.keyTransferMemSize = memInfos[TRANSFER_MEM_INFO_KEY_IDX].size();
    para.valueTransferMem = memInfos[TRANSFER_MEM_INFO_VALUE_IDX].ptr();
    para.valueTransferMemSize = memInfos[TRANSFER_MEM_INFO_VALUE_IDX].size();

    para.rdmaEnveInfosTransferMem = memInfos[TRANSFER_MEM_INFO_RDMA_ENVELOPE_IDX].ptr();
    para.rdmaEnveInfosTransferMemSize = memInfos[TRANSFER_MEM_INFO_RDMA_ENVELOPE_IDX].size();

    para.workerspaceMemSize = para.keyTransferMemSize + para.valueTransferMemSize +
        para.rdmaEnveInfosTransferMemSize;

    HCCL_INFO("[Remote][update]keyTransferMemSize[%llu] valueTransferMemSize[%llu] rdmaEnveInfosTransferMemSize[%llu]"
        " workerspaceMemSize[%llu] keyMaxNum[%llu] valueItemSize[%d]", para.keyTransferMemSize,
        para.valueTransferMemSize, para.rdmaEnveInfosTransferMemSize, para.workerspaceMemSize,
        ptr->param_.keyMaxNum, embeddingParam.valueItemSize);

    para.haveRdmaConn = ptr->param_.haveRdmaConn;
    para.tableId = ptr->param_.tableId;
    para.tableIdAddr = ptr->param_.tableIdAddr;
    para.pairedMode = ptr->param_.pairedMode;
    para.globalStepAddr = ptr->param_.globalStepAddr;

    if (para.pairedMode) { // paired类型
        para.indices = ptr->param_.indices;
        para.numUniqued = ptr->param_.numUniqued;
        para.psSeg = ptr->param_.psSeg;
        para.psSegNum = ptr->param_.psSegNum;
    }

    CHK_RET(CollRemoteUpdateKernelLaunch(para, embeddingParam.stream));

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::CollRemoteUpdatePaired(HcclOpPtr opPtr)
{
    return CollRemoteUpdate(opPtr);
}

HcclResult HcclImplRpc::WaitLookUpData(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    auto &alg = tagAlgMap_[ptr->param_.tag];
    CHK_SMART_PTR_NULL(alg);
    CHK_RET(alg->Process(opPtr));
    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::GetLookupRequest(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    ptr->param_.userRank = userRank_;

    EXECEPTION_CATCH({ auto &alg = GenerateEmbAlg<AlgRemoteLookUp>(ptr->param_.tag); \
        CHK_SMART_PTR_NULL(alg); \
        CHK_RET(alg->Process(opPtr)); \
        }, return HCCL_E_PTR);

    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::IsetLookupResponse(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    auto it = tagRankTransMemsMap_.Find(ptr->param_.tag);
    if (it.second) {
        ptr->param_.getRequestParam.transferMems = it.first->second;
    }

    auto &alg = tagAlgMap_[ptr->param_.tag];
    CHK_SMART_PTR_NULL(alg);
    CHK_RET(alg->Process(opPtr));
    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::RemoteLookUp(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);
    ptr->param_.userRank = userRank_;

    EXECEPTION_CATCH({ auto &alg = GenerateEmbAlg<AlgRemoteLookUp>(ptr->param_.tag); \
        CHK_SMART_PTR_NULL(alg); \
        CHK_RET(alg->Process(opPtr)); \
        }, return HCCL_E_PTR);

    CHK_RET(callbackTask_->CallbackRegStream(ptr->param_.embeddingParam.stream));
    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::IsendUpdateResponse(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);

    auto it = tagRankTransMemsMap_.Find(ptr->param_.tag);
    if (it.second) {
        ptr->param_.getRequestParam.transferMems = it.first->second;
    }

    auto &alg = tagAlgMap_[ptr->param_.tag];
    CHK_SMART_PTR_NULL(alg);

    CHK_RET(alg->Process(opPtr));
    CHK_RET(tagOpExecutorMap_[ptr->param_.tag]->ExecuteOp(opPtr));
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::CollRemoteUpdateKernelLaunch(HdcsEmbeddingServiceParam &para, void *stream)
{
    uint64_t beginTime = hrtMsprofSysCycleTime();
    rtKernelLaunchNames_t launchNames;
    launchNames.soName = "libccl_kernel.so";

    auto &profilingManager = hccl::ProfilingManager::Instance();

    // 下发ValueClear的aicore算子
    beginTime = hrtMsprofSysCycleTime();
    EsUpdatePara esUpdatePara;
    CHK_RET(SetupEsUpdatePara(esUpdatePara, para));
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
            ProfTaskType::TASK_VALUE_CLEAR_AICORE));
    CHK_RET(ValueClearAicoreKernelLaunch(para, stream));
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_VALUE_CLEAR_AICORE));
    CHK_RET(PrepareHostEnvelopeTemplates(para));

    beginTime = hrtMsprofSysCycleTime();
    // 使用输入的恢复矩阵或者自己生成的恢复矩阵进行key去重
    launchNames.opName = "CollRemoteUpdateKeyReduce";
    launchNames.kernelName = "HdcsCollRemoteUpdateKeyReduce";
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
        ProfTaskType::TASK_REMOTE_UPDATE_KEY_REDUCE));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsEmbeddingServiceParam), nullptr, stream));
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_REMOTE_UPDATE_KEY_REDUCE));

#ifdef USE_AICORE_REDUCESUM
    // 下发ReduceSum的aicore算子
    beginTime = hrtMsprofSysCycleTime();
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
        ProfTaskType::TASK_VALUE_REDUCE_SUM_AICORE));
    CHK_RET(ValueReduceSumAicoreKernelLaunch(para, stream));
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime, ProfTaskType::TASK_VALUE_REDUCE_SUM_AICORE));
#endif

    CHK_RET(AsyncCopyHostEnvelopeToDev(para, stream));

    launchNames.opName = "CollRemoteUpdateSendRequest";
    launchNames.kernelName = "HdcsCollRemoteUpdateSendRequest";
    beginTime = hrtMsprofSysCycleTime();
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
        ProfTaskType::TASK_REMOTE_UPDATE_SEND_REQUEST));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsEmbeddingServiceParam), nullptr, stream));

    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime,
        ProfTaskType::TASK_REMOTE_UPDATE_SEND_REQUEST));

    // rdma发送编排要在sendRequest之后
    CHK_RET(UpdateTaskSendByRdma(para, stream));

    // 下发notify和reset操作
    CHK_RET(NotifyAndResetKernelLaunch(para, stream));

    // 每次发送完成之后自动清理去重mappingInfo，不需要单独下发去重信息清理逻辑
    launchNames.opName = "CollRemoteUpdateRecvResponse";
    launchNames.kernelName = "HdcsCollRemoteUpdateRecvResponse";
    beginTime = hrtMsprofSysCycleTime();
    CHK_RET(profilingManager.CallMsprofReportAdditionInfoForEsUpdate(esUpdatePara,
        ProfTaskType::TASK_REMOTE_UPDATE_RECV_RESPONSE));
    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &para, sizeof(HdcsEmbeddingServiceParam), nullptr, stream));
    CHK_RET(profilingManager.CallEsMsprofReportTaskApi(false, beginTime,
        ProfTaskType::TASK_REMOTE_UPDATE_RECV_RESPONSE));

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::RegTransportsKernelLaunch(s32 tag, const map<u32, u32> &psRankIdMap)
{
    u32 rdmaTransportsNum = 0;
    HdcsRegTransportPara param{};

    for (auto &it : psRankIdMap) {
        HdcsTransportInfo &transportInfo = param.transportInfos[it.first];
        TransportHeterog *transportHeterog = transportHeterogMap_[tag][it.second].get();
        CHK_PTR_NULL(transportHeterog);

        TransportShmEvent *transport = dynamic_cast<TransportShmEvent *>(transportHeterog);
        if (transport == nullptr) {
            TransportHeterogRoce *transportRoce = dynamic_cast<TransportHeterogRoce *>(transportHeterog);
            CHK_PRT_RET(transportRoce == nullptr,
                HCCL_ERROR("tag[%u] transport type is not TransportShmEvent or TransportHeterogRoce.",
                tag), HCCL_E_PARA);

            transportInfo.commType = HcclHeterogCommType::RDMA;

            HCCL_INFO("tag[%d], rank[%u], psId[%u] transport type is TransportHeterogRoce", tag, it.second, it.first);
            rdmaTransportsNum++;
            continue;
        }

        TransportShmEventMember transportMember;
        CHK_RET(transport->GetTransportMember(transportMember));

        param.pcieTransportNum++;
        transportInfo.commType = HcclHeterogCommType::PCIE;
        transportInfo.tag = tag;
        transportInfo.rank = it.second;
        transportInfo.deviceLogicId = transportMember.deviceLogicId;
        transportInfo.inputShmMem = transportMember.inputShmMem.ptr();
        transportInfo.inputMemSize = transportMember.inputShmMem.size();
        transportInfo.outputShmMem = transportMember.outputShmMem.ptr();
        transportInfo.outputMemSize = transportMember.outputShmMem.size();
        transportInfo.envelopMem = transportMember.envelopeShmQue.ptr();
        transportInfo.envelopeSize = transportMember.envelopeShmQue.size();

        HCCL_INFO("tag[%d], rank[%u], psId[%u] transport type is TransportShmEvent", tag, it.second, it.first);
    }

    param.transportNum = psRankIdMap.size();

    HCCL_INFO("There are %d ShmEvent Transports, and %u HeterogRoce Transports, and %d total Transports",
        param.pcieTransportNum, rdmaTransportsNum, param.transportNum);

    if (param.pcieTransportNum == 0) {
        HCCL_INFO("There is no ShmEventTransport need to reg to device.");
        return HCCL_SUCCESS;
    }

    Stream stream(StreamType::STREAM_TYPE_ONLINE);
    rtKernelLaunchNames_t launchNames;
    launchNames.kernelName = "HdcsRegTransport";
    launchNames.soName = "libccl_kernel.so";
    launchNames.opName = "RegTransport";

    CHK_RET(hrtAicpuKernelLaunch(launchNames, 1, &param, sizeof(HdcsRegTransportPara), nullptr, stream.ptr()));
    CHK_RET(hcclStreamSynchronize(stream.ptr()));

    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::AbortSelf(s32 tag)
{
    EmbeddingServiceParam param;
    param.opType = EmbeddingServiceParam::OperatorType::ABORT_SELF;
    HcclOperator<EmbeddingServiceParam> op(param);

    auto it = tagOpExecutorMap_.Find(tag);
    if (it.second) {
        auto &exePtr = it.first->second;
        exePtr->ExecuteOp(&op);
    }

    unique_lock<mutex> lock(tagServiceHandleMapMutex_);
    auto serviceHandleIter = tagServiceHandleMap_.find(tag);
    if (serviceHandleIter != tagServiceHandleMap_.end()) {
        if (serviceHandleIter->second != nullptr) {
            HcclOperator<EmbeddingServiceParam> *ptr =
                static_cast<HcclOperator<EmbeddingServiceParam> *>(serviceHandleIter->second);
            AlgRemoteLookUp *alg = static_cast<AlgRemoteLookUp *>(ptr->param_.algPtr);
            OpRemoteLookUpExecutor *opExecutorPtr = static_cast<OpRemoteLookUpExecutor *>(ptr->param_.opExecutorPtr);
            opExecutorPtr->ExecuteOp(&op);
            delete alg;
            ptr->param_.algPtr = nullptr;
            delete opExecutorPtr;
            ptr->param_.opExecutorPtr = nullptr;
            tagServiceHandleMap_.erase(tag);
        }
    }
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::RegTransportLinks(s32 linkNum, void *transportPara)
{
    TransportShmEventMember *para = reinterpret_cast<TransportShmEventMember *>(transportPara);
    for (int i = 0; i < linkNum; i++) {
        TransportEndPointInfo commRankTagKey(0, (para + i)->rank, (para + i)->tag);
        std::unique_lock<SpinMutex> transportMapLock(transportMapSpinMutex_);
        std::unique_ptr<TransportHeterog>& transportInfo = transportStorage_[commRankTagKey];
        HCCL_DEBUG("RegTransportLinks: commId[0] peerRank[%u] tag[%u]", (para + i)->rank, (para + i)->tag);
        transportMapLock.unlock();

        if (transportInfo == nullptr) {
            std::unique_ptr<TransportHeterog> transportPtr;
            MachinePara machinePara;
            CHK_RET(InitRecvMsgAndRequestBuffer());
            std::chrono::milliseconds kdefaultTimeout = std::chrono::seconds(GetExternalInputHcclLinkTimeOut());
            HcclIpAddress invalidIP;
            transportPtr.reset(new (std::nothrow) TransportShmEvent(nullptr,  notifyPool_, machinePara,
                kdefaultTimeout, invalidIP, invalidIP, 0, 0, (para + i)->deviceLogicId, CLIENT_ROLE_SOCKET,
                transportResourceInfo_, nullptr));
            CHK_SMART_PTR_NULL(transportPtr);
            TransportShmEvent *transport = dynamic_cast<TransportShmEvent *>(transportPtr.get());
            CHK_RET(transport->Init(*(para + i)));
            transportInfo = std::move(transportPtr);
        } else {
            TransportShmEvent *transport = dynamic_cast<TransportShmEvent *>(transportInfo.get());
            CHK_RET(transport->Init(*(para + i)));
        }

        HCCL_INFO("RegTransportLink rank[%u] tag[%u].", (para + i)->rank, (para + i)->tag);
    }
    return HCCL_SUCCESS;
}

bool HcclImplRpc::PsGetBoardVersion()
{
    return GetBoardVersion(clientsInfo_);
}

bool HcclImplRpc::WorkerGetBoardVersion()
{
    return GetBoardVersion(serversInfo_);
}

HcclResult HcclImplRpc::WaitSome(HcclOpPtr opPtr)
{
    HcclOperator<EmbeddingServiceParam> *ptr = static_cast<HcclOperator<EmbeddingServiceParam> *>(opPtr);

    *(ptr->param_.waitSomeParam.compCount) = 0;
    for (s32 i = 0; i < ptr->param_.waitSomeParam.requestCount; i++) {
        void *request = ptr->param_.waitSomeParam.requestArray[i];
        s32 tag = static_cast<HcclRequestInfo *>(request)->tag;
        auto& alg = tagAlgMap_[tag];
        if (alg == nullptr) {
            s32 flag = HCCL_TEST_INCOMPLETED;
            HcclStatus status;
            while (flag != HCCL_TEST_COMPLETED) {
                CHK_RET(HcclTest(request, flag, status));
                if (flag != HCCL_TEST_COMPLETED) {
                    SaluSleep(LOOP_SLEEP_TIME_US);
                }
            }
            ptr->param_.waitSomeParam.compIndices[*(ptr->param_.waitSomeParam.compCount)] = i;
            ptr->param_.waitSomeParam.compStatus[*(ptr->param_.waitSomeParam.compCount)] = status;
            (*(ptr->param_.waitSomeParam.compCount))++;
            continue;
        }
        CHK_RET(alg->Process(opPtr));
        auto &exe = tagOpExecutorMap_[tag];
        CHK_SMART_PTR_NULL(exe);
        CHK_RET(exe->ExecuteOp(opPtr));
    }
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::Isend(void *buffer, u64 count, HcclDataType dataType, u32 peerRank, s32 tag,
    HcclRequest &requestHandle, HcclUserRequire &userRequire)
{
    if ((buffer == nullptr) && (count != 0)) {
        HCCL_ERROR("[Check][Buffer]errNo[0x%016llx] buffer[%p] or count[%llu] is invalid",
            HCCL_ERROR_CODE(HCCL_E_PARA), buffer, count);
        return HCCL_E_PARA;
    }

    if (peerRank >= userRankSize_) {
        HCCL_ERROR("[Check][UserRank]errNo[0x%016llx] peerRank:[%u] is out of range[0 ~ %u]",
            HCCL_ERROR_CODE(HCCL_E_PARA), peerRank, userRankSize_);
        return HCCL_E_PARA;
    }

    TransportHandle transportHandle = nullptr;
    CHK_RET(BuildHeterogeneousTransport(0, peerRank, tag, transportHandle));

    TransportHeterog *transportPtr = reinterpret_cast<TransportHeterog *>(transportHandle);
    HcclRequestInfo* request = nullptr;
    TransData sendData(reinterpret_cast<u64>(buffer), reinterpret_cast<u64>(nullptr), count, dataType, false,
        userRequire.tableId, userRequire.globalStep);
    TransportEndPointInfo srcEp(0, userRank_, tag);
    TransportEndPointInfo dstEp(0, peerRank, tag);
    TransportEndPointParam epParam(srcEp, dstEp);
    CHK_RET(transportPtr->Isend(sendData, epParam, request));
    request->commHandle = commHandle_;
    requestHandle = request;
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::BuildHeterogeneousTransport(u32 commId, u32 peerRank, s32 tag,
    TransportHandle &transportHandle)
{
    TransportEndPointInfo commRankTagKey(commId, peerRank, tag);
    std::unique_lock<SpinMutex> transportMapLock(transportMapSpinMutex_);
    std::unique_ptr<TransportHeterog>& transportInfo = transportStorage_[commRankTagKey];
    transportMapLock.unlock();

    if (transportInfo == nullptr) {
        std::string transTag;
        if (userRank_ > peerRank) {
            transTag = collectiveId_ + "_" + std::to_string(peerRank) + "_" + std::to_string(userRank_) + "_";
        } else {
            transTag = collectiveId_ + "_" + std::to_string(userRank_) + "_" + std::to_string(peerRank) + "_";
        }
        transTag += std::to_string(tag);
        std::unique_ptr<TransportHeterog> transportPtr;
        if (GetExternalInputHcclIsTcpMode()) {
            transportPtr.reset(new (std::nothrow) TransportHeterogEventTcp(transTag, rankInfoList_[userRank_].nicIp[0],
            rankInfoList_[peerRank].nicIp[0], ranksPort_[peerRank], ranksPort_[userRank_], devicePhyId_,
            transportResourceInfo_));
        } else {
            transportPtr.reset(new (std::nothrow) TransportHeterogEventRoce(transTag, rankInfoList_[userRank_].nicIp[0],
            rankInfoList_[peerRank].nicIp[0], ranksPort_[peerRank], ranksPort_[userRank_], transportResourceInfo_));
        }
        CHK_SMART_PTR_NULL(transportPtr);
        CHK_RET(transportPtr->SetDeviceIndex(deviceLogicId_));
        CHK_RET(transportPtr->Init());
        transportInfo = std::move(transportPtr);
    }
    transportHandle = transportInfo.get();
    CHK_PTR_NULL(transportHandle);
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::Improbe(u32 peerRank, s32 tag, s32 &flag, HcclMessage &msgHandle, HcclStatus &status)
{
    if (peerRank >= userRankSize_) {
        HCCL_ERROR("[Check][UserRank]errNo[0x%016llx] peerRank:[%u] is out of range[0 ~ %u]",
            HCCL_ERROR_CODE(HCCL_E_PARA), peerRank, userRankSize_);
        return HCCL_E_PARA;
    }

    void* transportHandle = nullptr;
    CHK_RET(BuildHeterogeneousTransport(0, peerRank, tag, transportHandle));

    TransportHeterog *transportPtr = reinterpret_cast<TransportHeterog *>(transportHandle);
    TransportEndPointInfo srcEp(0, peerRank, tag);
    TransportEndPointInfo dstEp(0, userRank_, tag);
    TransportEndPointParam epParam(srcEp, dstEp);
    HcclMessageInfo *msg = nullptr;
    CHK_RET(transportPtr->Improbe(epParam, flag, msg, status));
    msgHandle = msg;
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::Imrecv(void* buffer, u64 count, HcclDataType dataType, HcclMessage msgHandle,
    HcclRequest &requestHandle)
{
    HcclMessageInfo* msg = static_cast<HcclMessageInfo *>(msgHandle);
    CHK_PTR_NULL(msg);
    TransportHeterog *transportPtr = reinterpret_cast<TransportHeterog *>(msg->transportHandle);
    CHK_PTR_NULL(transportPtr);

    HcclRequestInfo* request = nullptr;
    TransData recvData(reinterpret_cast<u64>(nullptr), reinterpret_cast<u64>(buffer), count, dataType);
    CHK_RET(transportPtr->Imrecv(recvData, *msg, request));
    requestHandle = request;
    return HCCL_SUCCESS;
}

HcclResult HcclImplRpc::HcclTest(HcclRequest requestHandle, s32 &flag, HcclStatus &compState)
{
    HcclRequestInfo *request = reinterpret_cast<HcclRequestInfo *>(requestHandle);
    CHK_PTR_NULL(request->transportHandle);
    TransportHeterog *transportPtr = reinterpret_cast<TransportHeterog *>(request->transportHandle);
    return transportPtr->Test(*request, flag, compState);
}

HcclResult HcclImplRpc::SetupEsUpdatePara(EsUpdatePara &esUpdatePara, const HdcsEmbeddingServiceParam &para)
{
    esUpdatePara.tag = para.tag;
    esUpdatePara.srcAddr = para.keys;
    esUpdatePara.dstAddr = para.values;
    esUpdatePara.dataSize = (u32)para.keyMaxNum;
    CHK_SAFETY_FUNC_RET(memcpy_s(esUpdatePara.groupName, sizeof(esUpdatePara.groupName), para.groupName,
        sizeof(para.groupName)));
    return HCCL_SUCCESS;
}

}